package com.wt.admin.service.ai.impl.agents.easyai;

import com.fasterxml.jackson.core.type.TypeReference;
import com.wt.admin.service.ai.impl.agents.easyai.api.EasyAIApi;
import com.wt.admin.service.ai.impl.agents.easyai.api.EasyAIOptions;
import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.DefaultUsage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.ai.chat.observation.ChatModelObservationContext;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.tool.*;
import org.springframework.ai.ollama.api.common.OllamaApiConstants;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.util.json.JsonParser;
import org.springframework.retry.support.RetryTemplate;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

import java.time.Duration;
import java.util.*;

public class EasyAIChatModel implements ChatModel {

    private static final Logger logger = LoggerFactory.getLogger(EasyAIChatModel.class);
    private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
    private static final ToolCallingManager DEFAULT_TOOL_CALLING_MANAGER = ToolCallingManager.builder().build();
    private final EasyAIApi easyAIApi;
    private final EasyAIOptions defaultOptions;
    private final ObservationRegistry observationRegistry;
    private final ToolCallingManager toolCallingManager;
    private final ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate;
    private ChatModelObservationConvention observationConvention;
    private final RetryTemplate retryTemplate;

    public EasyAIChatModel(EasyAIApi easyAIApi, EasyAIOptions defaultOptions, ToolCallingManager toolCallingManager, ObservationRegistry observationRegistry, ModelManagementOptions modelManagementOptions, ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate, RetryTemplate retryTemplate) {
        this.observationConvention = DEFAULT_OBSERVATION_CONVENTION;
        Assert.notNull(easyAIApi, "easyAIApi must not be null");
        Assert.notNull(defaultOptions, "easyAIOptions must not be null");
        Assert.notNull(toolCallingManager, "toolCallingManager must not be null");
        Assert.notNull(observationRegistry, "observationRegistry must not be null");
        Assert.notNull(modelManagementOptions, "modelManagementOptions must not be null");
        Assert.notNull(toolExecutionEligibilityPredicate, "toolExecutionEligibilityPredicate must not be null");
        Assert.notNull(retryTemplate, "retryTemplate must not be null");
        this.easyAIApi = easyAIApi;
        this.defaultOptions = defaultOptions;
        this.toolCallingManager = toolCallingManager;
        this.observationRegistry = observationRegistry;
        this.toolExecutionEligibilityPredicate = toolExecutionEligibilityPredicate;
        this.retryTemplate = retryTemplate;
    }

    @Deprecated
    public EasyAIChatModel(EasyAIApi easyAIApi, EasyAIOptions defaultOptions,ToolCallingManager toolCallingManager,  ObservationRegistry observationRegistry, ModelManagementOptions modelManagementOptions) {
        this(easyAIApi, defaultOptions, toolCallingManager, observationRegistry, modelManagementOptions, new DefaultToolExecutionEligibilityPredicate(), RetryUtils.DEFAULT_RETRY_TEMPLATE);
    }

    @Override
    public ChatResponse call(Prompt prompt) {
        Prompt requestPrompt = this.buildRequestPrompt(prompt);
        return this.internalCall(requestPrompt, null);
    }

    @Override
    public Flux<ChatResponse> stream(Prompt prompt) {
        Prompt requestPrompt = buildRequestPrompt(prompt);
        return this.internalStream(requestPrompt, null);
    }

    private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatResponse) {
        EasyAIApi.ChatRequest request = this.chatRequest(prompt, false);
        ChatModelObservationContext observationContext = ChatModelObservationContext.builder().prompt(prompt).provider(OllamaApiConstants.PROVIDER_NAME).build();
        ChatResponse response = (ChatResponse)ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, this.observationRegistry).observe(() -> {
            EasyAIApi.ChatResponse ollamaResponse = (EasyAIApi.ChatResponse)this.retryTemplate.execute((ctx) -> this.easyAIApi.chat(request));
            List<AssistantMessage.ToolCall> toolCalls = ollamaResponse.message().toolCalls() == null ? List.of() : ollamaResponse.message().toolCalls().stream().map((toolCall) -> new AssistantMessage.ToolCall("", "function", toolCall.function().name(), ModelOptionsUtils.toJsonString(toolCall.function().arguments()))).toList();
            AssistantMessage assistantMessage = new AssistantMessage(ollamaResponse.message().content(), Map.of(), toolCalls);
            ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
            if (ollamaResponse.promptEvalCount() != null && ollamaResponse.evalCount() != null) {
                generationMetadata = ChatGenerationMetadata.builder().finishReason(ollamaResponse.doneReason()).build();
            }

            Generation generator = new Generation(assistantMessage, generationMetadata);
            ChatResponse chatResponse = new ChatResponse(List.of(generator), from(ollamaResponse, previousChatResponse));
            observationContext.setResponse(chatResponse);
            return chatResponse;
        });
        if (this.toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response)) {
            ToolExecutionResult toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response);
            return toolExecutionResult.returnDirect() ? ChatResponse.builder().from(response).generations(ToolExecutionResult.buildGenerations(toolExecutionResult)).build() : this.internalCall(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), response);
        } else {
            return response;
        }
    }

    private Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousChatResponse) {
        return Flux.deferContextual((contextView) -> {
            EasyAIApi.ChatRequest request = this.chatRequest(prompt, true);
            ChatModelObservationContext observationContext = ChatModelObservationContext.builder().prompt(prompt).provider(OllamaApiConstants.PROVIDER_NAME).build();
            Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, this.observationRegistry);
            observation.parentObservation((Observation)contextView.getOrDefault("micrometer.observation", (Object)null)).start();
            Flux<EasyAIApi.ChatResponse> easyaiResponse = this.easyAIApi.streamingChat(request);
            Flux<ChatResponse> chatResponse = easyaiResponse.map((chunk) -> {
                String content = chunk.message() != null ? chunk.message().content() : "";
                List<AssistantMessage.ToolCall> toolCalls = List.of();
                if (chunk.message() != null && chunk.message().toolCalls() != null) {
                    toolCalls = chunk.message().toolCalls().stream().map((toolCall) -> new AssistantMessage.ToolCall("", "function", toolCall.function().name(), ModelOptionsUtils.toJsonString(toolCall.function().arguments()))).toList();
                }

                AssistantMessage assistantMessage = new AssistantMessage(content, Map.of(), toolCalls);
                ChatGenerationMetadata generationMetadata = ChatGenerationMetadata.NULL;
                if (chunk.promptEvalCount() != null && chunk.evalCount() != null) {
                    generationMetadata = ChatGenerationMetadata.builder().finishReason(chunk.doneReason()).build();
                }

                Generation generator = new Generation(assistantMessage, generationMetadata);
                return new ChatResponse(List.of(generator), from(chunk, previousChatResponse));
            });
            Flux<ChatResponse> var10000 = chatResponse.flatMap((response) -> this.toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response) ? Flux.defer(() -> {
                ToolExecutionResult toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response);
                return toolExecutionResult.returnDirect() ? Flux.just(ChatResponse.builder().from(response).generations(ToolExecutionResult.buildGenerations(toolExecutionResult)).build()) : this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), response);
            }).subscribeOn(Schedulers.boundedElastic()) : Flux.just(response));
            Objects.requireNonNull(observation);
            Flux<ChatResponse> chatResponseFlux = var10000.doOnError(observation::error).doFinally((s) -> observation.stop()).contextWrite((ctx) -> ctx.put("micrometer.observation", observation));
            MessageAggregator var10 = new MessageAggregator();
            Objects.requireNonNull(observationContext);
            return var10.aggregate(chatResponseFlux, observationContext::setResponse);
        });
    }

    EasyAIApi.ChatRequest chatRequest(Prompt prompt, boolean stream) {

        List<EasyAIApi.Message> ollamaMessages = prompt.getInstructions().stream().map(message -> {
            if (message instanceof UserMessage userMessage) {
                var messageBuilder = EasyAIApi.Message.builder(EasyAIApi.Message.Role.USER).content(message.getText());
                if (!CollectionUtils.isEmpty(userMessage.getMedia())) {
                    messageBuilder.images(
                            userMessage.getMedia().stream().map(media -> this.fromMediaData(media.getData())).toList());
                }
                return List.of(messageBuilder.build());
            }
            else if (message instanceof SystemMessage systemMessage) {
                return List.of(EasyAIApi.Message.builder(EasyAIApi.Message.Role.SYSTEM).content(systemMessage.getText()).build());
            }
            else if (message instanceof AssistantMessage assistantMessage) {
                List<EasyAIApi.Message.ToolCall> toolCalls = null;
                if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
                    toolCalls = assistantMessage.getToolCalls().stream().map(toolCall -> {
                        var function = new EasyAIApi.Message.ToolCallFunction(toolCall.name(),
                                JsonParser.fromJson(toolCall.arguments(), new TypeReference<>() {
                                }));
                        return new EasyAIApi.Message.ToolCall(function);
                    }).toList();
                }
                return List.of(EasyAIApi.Message.builder(EasyAIApi.Message.Role.ASSISTANT)
                        .content(assistantMessage.getText())
                        .toolCalls(toolCalls)
                        .build());
            }
            else if (message instanceof ToolResponseMessage toolMessage) {
                return toolMessage.getResponses()
                        .stream()
                        .map(tr -> EasyAIApi.Message.builder(EasyAIApi.Message.Role.TOOL).content(tr.responseData()).build())
                        .toList();
            }
            throw new IllegalArgumentException("Unsupported message type: " + message.getMessageType());
        }).flatMap(List::stream).toList();

        EasyAIOptions requestOptions =  (EasyAIOptions) prompt.getOptions();

        EasyAIApi.ChatRequest.Builder requestBuilder = EasyAIApi.ChatRequest.builder(requestOptions.getModel())
                .stream(stream)
                .messages(ollamaMessages)
                .options(requestOptions);

        if (requestOptions.getFormat() != null) {
            requestBuilder.format(requestOptions.getFormat());
        }

        if (requestOptions.getKeepAlive() != null) {
            requestBuilder.keepAlive(requestOptions.getKeepAlive());
        }

        List<ToolDefinition> toolDefinitions = this.toolCallingManager.resolveToolDefinitions(requestOptions);
        if (!CollectionUtils.isEmpty(toolDefinitions)) {
            requestBuilder.tools(this.getTools(toolDefinitions));
        }

        return requestBuilder.build();
    }

    EasyAIApi.ChatRequest easyAIChatRequest(Prompt prompt, boolean stream) {
        List<EasyAIApi.Message> easyAIMessages = prompt.getInstructions().stream().map((message) -> {
            if (message instanceof UserMessage userMessage) {
                EasyAIApi.Message.Builder messageBuilder = EasyAIApi.Message.builder(EasyAIApi.Message.Role.USER).content(message.getText());
                if (!CollectionUtils.isEmpty(userMessage.getMedia())) {
                    messageBuilder.images(userMessage.getMedia().stream().map((media) -> {
                        return this.fromMediaData(media.getData());
                    }).toList());
                }

                return List.of(messageBuilder.build());
            } else if (message instanceof SystemMessage systemMessage) {
                return List.of(EasyAIApi.Message.builder(EasyAIApi.Message.Role.SYSTEM).content(systemMessage.getText()).build());
            } else if (message instanceof AssistantMessage assistantMessage) {
                List<EasyAIApi.Message.ToolCall> toolCalls = null;
                if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
                    toolCalls = assistantMessage.getToolCalls().stream().map((toolCall) -> {
                        EasyAIApi.Message.ToolCallFunction function = new EasyAIApi.Message.ToolCallFunction(toolCall.name(), JsonParser.fromJson(toolCall.arguments(), new TypeReference<Map<String, Object>>() {
                        }));
                        return new EasyAIApi.Message.ToolCall(function);
                    }).toList();
                }

                return List.of(EasyAIApi.Message.builder(EasyAIApi.Message.Role.ASSISTANT).content(assistantMessage.getText()).toolCalls(toolCalls).build());
            } else if (message instanceof ToolResponseMessage toolMessage) {
                return toolMessage.getResponses().stream().map((tr) -> {
                    return EasyAIApi.Message.builder(EasyAIApi.Message.Role.TOOL).content(tr.responseData()).build();
                }).toList();
            } else {
                throw new IllegalArgumentException("Unsupported message type: " + String.valueOf(message.getMessageType()));
            }
        }).flatMap(Collection::stream).toList();
        EasyAIOptions requestOptions = (EasyAIOptions) prompt.getOptions();
        EasyAIApi.ChatRequest.Builder requestBuilder = EasyAIApi.ChatRequest.builder(requestOptions.getModel())
                .stream(stream).messages(easyAIMessages)
                .options(requestOptions);
        if (requestOptions.getFormat() != null) {
            requestBuilder.format(requestOptions.getFormat());
        }

        if (requestOptions.getKeepAlive() != null) {
            requestBuilder.keepAlive(requestOptions.getKeepAlive());
        }

        List<ToolDefinition> toolDefinitions = this.toolCallingManager.resolveToolDefinitions(requestOptions);
        if (!CollectionUtils.isEmpty(toolDefinitions)) {
            requestBuilder.tools(this.getTools(toolDefinitions));
        }
        return requestBuilder.build();
    }

    private String fromMediaData(Object mediaData) {
        if (mediaData instanceof byte[] bytes) {
            return Base64.getEncoder().encodeToString(bytes);
        } else if (mediaData instanceof String text) {
            return text;
        } else {
            throw new IllegalArgumentException("Unsupported media data type: " + mediaData.getClass().getSimpleName());
        }
    }

    private List<EasyAIApi.ChatRequest.Tool> getTools(List<ToolDefinition> toolDefinitions) {
        return toolDefinitions.stream().map((toolDefinition) -> {
            EasyAIApi.ChatRequest.Tool.Function tool = new EasyAIApi.ChatRequest.Tool.Function(toolDefinition.name(), toolDefinition.description(), toolDefinition.inputSchema());
            return new EasyAIApi.ChatRequest.Tool(tool);
        }).toList();
    }

    Prompt buildRequestPrompt(Prompt prompt) {
        EasyAIOptions runtimeOptions = null;
        if (prompt.getOptions() != null) {
            ChatOptions var5 = prompt.getOptions();
            if (var5 instanceof ToolCallingChatOptions) {
                ToolCallingChatOptions toolCallingChatOptions = (ToolCallingChatOptions) var5;
                runtimeOptions = ModelOptionsUtils.copyToTarget(toolCallingChatOptions, ToolCallingChatOptions.class, EasyAIOptions.class);
            } else {
                runtimeOptions = ModelOptionsUtils.copyToTarget(prompt.getOptions(), ChatOptions.class, EasyAIOptions.class);
            }
        }
        EasyAIOptions requestOptions = ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions, EasyAIOptions.class);
        if (runtimeOptions != null) {
            requestOptions.setInternalToolExecutionEnabled((Boolean) ModelOptionsUtils.mergeOption(runtimeOptions.getInternalToolExecutionEnabled(), this.defaultOptions.getInternalToolExecutionEnabled()));
            requestOptions.setToolNames(ToolCallingChatOptions.mergeToolNames(runtimeOptions.getToolNames(), this.defaultOptions.getToolNames()));
            requestOptions.setToolCallbacks(ToolCallingChatOptions.mergeToolCallbacks(runtimeOptions.getToolCallbacks(), this.defaultOptions.getToolCallbacks()));
            requestOptions.setToolContext(ToolCallingChatOptions.mergeToolContext(runtimeOptions.getToolContext(), this.defaultOptions.getToolContext()));
        } else {
            requestOptions.setInternalToolExecutionEnabled(this.defaultOptions.getInternalToolExecutionEnabled());
            requestOptions.setToolNames(this.defaultOptions.getToolNames());
            requestOptions.setToolCallbacks(this.defaultOptions.getToolCallbacks());
            requestOptions.setToolContext(this.defaultOptions.getToolContext());
        }
        if (!StringUtils.hasText(requestOptions.getModel())) {
            throw new IllegalArgumentException("model cannot be null or empty");
        } else {
            ToolCallingChatOptions.validateToolCallbacks(requestOptions.getToolCallbacks());
            return new Prompt(prompt.getInstructions(), requestOptions);
        }
    }


    static ChatResponseMetadata from(EasyAIApi.ChatResponse response, ChatResponse previousChatResponse) {
        Assert.notNull(response, "EasyAIApi.ChatResponse must not be null");
        DefaultUsage newUsage = getDefaultUsage(response);
        Integer promptTokens = newUsage.getPromptTokens();
        Integer generationTokens = newUsage.getCompletionTokens();
        int totalTokens = newUsage.getTotalTokens();
        Duration evalDuration = response.getEvalDuration();
        Duration promptEvalDuration = response.getPromptEvalDuration();
        Duration loadDuration = response.getLoadDuration();
        Duration totalDuration = response.getTotalDuration();
        if (previousChatResponse != null && previousChatResponse.getMetadata() != null) {
            if (previousChatResponse.getMetadata().get("eval-duration") != null) {
                evalDuration = evalDuration.plus((Duration) previousChatResponse.getMetadata().get("eval-duration"));
            }

            if (previousChatResponse.getMetadata().get("prompt-eval-duration") != null) {
                promptEvalDuration = promptEvalDuration.plus((Duration) previousChatResponse.getMetadata().get("prompt-eval-duration"));
            }

            if (previousChatResponse.getMetadata().get("load-duration") != null) {
                loadDuration = loadDuration.plus((Duration) previousChatResponse.getMetadata().get("load-duration"));
            }

            if (previousChatResponse.getMetadata().get("total-duration") != null) {
                totalDuration = totalDuration.plus((Duration) previousChatResponse.getMetadata().get("total-duration"));
            }

            if (previousChatResponse.getMetadata().getUsage() != null) {
                promptTokens = promptTokens + previousChatResponse.getMetadata().getUsage().getPromptTokens();
                generationTokens = generationTokens + previousChatResponse.getMetadata().getUsage().getCompletionTokens();
                totalTokens += previousChatResponse.getMetadata().getUsage().getTotalTokens();
            }
        }

        DefaultUsage aggregatedUsage = new DefaultUsage(promptTokens, generationTokens, totalTokens);
        return ChatResponseMetadata.builder().usage(aggregatedUsage).model(response.model()).keyValue("created-at", response.createdAt()).keyValue("eval-duration", evalDuration).keyValue("eval-count", aggregatedUsage.getCompletionTokens()).keyValue("load-duration", loadDuration).keyValue("prompt-eval-duration", promptEvalDuration).keyValue("prompt-eval-count", aggregatedUsage.getPromptTokens()).keyValue("total-duration", totalDuration).keyValue("done", response.done()).build();
    }

    private static DefaultUsage getDefaultUsage(EasyAIApi.ChatResponse response) {
        return new DefaultUsage((Integer) Optional.ofNullable(response.promptEvalCount()).orElse(0), (Integer) Optional.ofNullable(response.evalCount()).orElse(0));
    }

    public static Builder builder() {
        return new Builder();
    }

    public static final class Builder {
        private EasyAIApi easyAIAPI;
        private EasyAIOptions defaultOptions;
        private ToolCallingManager toolCallingManager;
        private ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate;
        private ObservationRegistry observationRegistry;
        private ModelManagementOptions modelManagementOptions;
        private RetryTemplate retryTemplate;


        private Builder() {
            this.observationRegistry = ObservationRegistry.NOOP;
            this.modelManagementOptions = ModelManagementOptions.defaults();
            this.retryTemplate = RetryUtils.DEFAULT_RETRY_TEMPLATE;
        }

        public Builder easyAIAPI(EasyAIApi easyAIAPI) {
            this.easyAIAPI = easyAIAPI;
            return this;
        }

        public Builder toolExecutionEligibilityPredicate(ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate) {
            this.toolExecutionEligibilityPredicate = toolExecutionEligibilityPredicate;
            return this;
        }

        public Builder options(EasyAIOptions defaultOptions) {
            this.defaultOptions = defaultOptions;
            return this;
        }

        public Builder toolCallingManager(ToolCallingManager toolCallingManager) {
            this.toolCallingManager = toolCallingManager;
            return this;
        }



        public Builder observationRegistry(ObservationRegistry observationRegistry) {
            this.observationRegistry = observationRegistry;
            return this;
        }

        public Builder modelManagementOptions(ModelManagementOptions modelManagementOptions) {
            this.modelManagementOptions = modelManagementOptions;
            return this;
        }

        public Builder retryTemplate(RetryTemplate retryTemplate) {
            this.retryTemplate = retryTemplate;
            return this;
        }

        public EasyAIChatModel build() {
            return this.toolCallingManager != null ? new EasyAIChatModel(this.easyAIAPI, this.defaultOptions, this.toolCallingManager, this.observationRegistry, this.modelManagementOptions, this.toolExecutionEligibilityPredicate,this.retryTemplate) : new EasyAIChatModel(this.easyAIAPI, this.defaultOptions, EasyAIChatModel.DEFAULT_TOOL_CALLING_MANAGER, this.observationRegistry, this.modelManagementOptions, this.toolExecutionEligibilityPredicate,this.retryTemplate);
        }
    }
}
